{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Check Token Usage\n",
    "\n",
    "- Author: [Haseom Shin](https://github.com/IHAGI-c)\n",
    "- Peer Review : [Teddy Lee](https://github.com/teddylee777), [Sooyoung](https://github.com/sooyoung-wind)\n",
    "- Proofread : [Two-Jay](https://github.com/Two-Jay)\n",
    "- This is a part of [LangChain Open Tutorial](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial)\n",
    "\n",
    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/04-Model/05-CheckTokenUsage.ipynb)[![Open in GitHub](https://img.shields.io/badge/Open%20in%20GitHub-181717?style=flat-square&logo=github&logoColor=white)](https://github.com/LangChain-OpenTutorial/LangChain-OpenTutorial/blob/main/04-Model/05-CheckTokenUsage.ipynb)\n",
    "## Overview\n",
    "\n",
    "This tutorial covers how to track and monitor token usage with ```LangChain``` and ```OpenAI API```.\n",
    "\n",
    "```Token usage tracking``` is crucial for managing API costs and optimizing resource utilization.\n",
    "\n",
    "In this tutorial, we will create a simple example to measure and monitor token consumption during OpenAI API calls using LangChain's ```CallbackHandler```.\n",
    "\n",
    "![example](./assets/04-CheckTokenUsage-example-flow-token-usage.png)\n",
    "\n",
    "### Table of Contents\n",
    "\n",
    "- [Overview](#overview)\n",
    "- [Environment Setup](#environment-setup)\n",
    "- [Implementing Check Token Usage](#implementing-check-token-usage)\n",
    "- [Monitoring Token Usage](#monitoring-token-usage)\n",
    "\n",
    "### References\n",
    "\n",
    "- [OpenAI API Pricing](https://openai.com/api/pricing/)\n",
    "- [Token Usage Guide](https://help.openai.com/en/articles/4936856-what-are-tokens-and-how-to-count-them)\n",
    "- [LangChain Python API Reference](https://python.langchain.com/api_reference/community/callbacks/langchain_community.callbacks.manager.get_openai_callback.html)\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Environment Setup\n",
    "\n",
    "Set up the environment. You may refer to [Environment Setup](https://wikidocs.net/257836) for more details.\n",
    "\n",
    "**[Note]**\n",
    "- ```langchain-opentutorial``` is a package that provides a set of easy-to-use environment setup, useful functions and utilities for tutorials. \n",
    "- You can checkout the [```langchain-opentutorial```](https://github.com/LangChain-OpenTutorial/langchain-opentutorial-pypi) for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%capture --no-stderr\n",
    "%pip install langchain-opentutorial"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m24.2\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.3.1\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# Install required packages\n",
    "from langchain_opentutorial import package\n",
    "\n",
    "package.install(\n",
    "    [\n",
    "        \"langsmith\",\n",
    "        \"langchain\",\n",
    "        \"langchain_openai\",\n",
    "        \"langchain_community\",\n",
    "    ],\n",
    "    verbose=False,\n",
    "    upgrade=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set environment variables\n",
    "from langchain_opentutorial import set_env\n",
    "\n",
    "set_env(\n",
    "    {\n",
    "        \"OPENAI_API_KEY\": \"\",\n",
    "        \"LANGCHAIN_API_KEY\": \"\",\n",
    "        \"LANGCHAIN_TRACING_V2\": \"true\",\n",
    "        \"LANGCHAIN_ENDPOINT\": \"https://api.smith.langchain.com\",\n",
    "        \"LANGCHAIN_PROJECT\": \"04-CheckTokenUsage\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can alternatively set ```OPENAI_API_KEY``` in ```.env``` file and load it. \n",
    "\n",
    "[Note] This is not necessary if you've already set ```OPENAI_API_KEY``` in previous steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's setup ```ChatOpenAI``` with ```gpt-4o``` model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "# Load the model\n",
    "llm = ChatOpenAI(model_name=\"gpt-4o\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Implementing Check Token Usage\n",
    "\n",
    "if you want to check token usage, you can use ```get_openai_callback``` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokens Used: 28\n",
      "\tPrompt Tokens: 15\n",
      "\t\tPrompt Tokens Cached: 0\n",
      "\tCompletion Tokens: 13\n",
      "\t\tReasoning Tokens: 0\n",
      "Successful Requests: 1\n",
      "Total Cost (USD): $0.00016749999999999998\n"
     ]
    }
   ],
   "source": [
    "# callback to track it\n",
    "from langchain_community.callbacks.manager import get_openai_callback\n",
    "\n",
    "with get_openai_callback() as cb:\n",
    "    result = llm.invoke(\"where is the capital of United States?\")\n",
    "    print(cb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total tokens used: \t\t28\n",
      "Tokens used in prompt: \t\t15\n",
      "Tokens used in completion: \t13\n",
      "Cost: \t\t\t\t$0.00016749999999999998\n"
     ]
    }
   ],
   "source": [
    "# callback to track it\n",
    "with get_openai_callback() as cb:\n",
    "    result = llm.invoke(\"where is the capital of United States?\")\n",
    "    print(f\"Total tokens used: \\t\\t{cb.total_tokens}\")\n",
    "    print(f\"Tokens used in prompt: \\t\\t{cb.prompt_tokens}\")\n",
    "    print(f\"Tokens used in completion: \\t{cb.completion_tokens}\")\n",
    "    print(f\"Cost: \\t\\t\\t\\t${cb.total_cost}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Monitoring Token Usage\n",
    "\n",
    "Token usage monitoring is crucial for managing costs and resources when using the OpenAI API. LangChain provides an easy way to track this through ```get_openai_callback()```.\n",
    "\n",
    "In this section, we'll explore token usage monitoring through three key scenarios:\n",
    "\n",
    "1. **Single Query Monitoring**: \n",
    "   - Track token usage for individual API calls\n",
    "   - Distinguish between prompt and completion tokens\n",
    "   - Calculate costs\n",
    "\n",
    "2. **Multiple Queries Monitoring**:\n",
    "   - Track cumulative token usage across multiple API calls\n",
    "   - Analyze total costs\n",
    "\n",
    "> **Note**: Token usage monitoring is currently only supported for OpenAI API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1. Single Query Monitoring\n",
      "----------------------------------------\n",
      "Response: The capital of France is Paris.\n",
      "----------------------------------------\n",
      "Token Usage Statistics:\n",
      "Prompt Tokens: \t\t14\n",
      "Completion Tokens: \t8\n",
      "Total Tokens: \t\t22\n",
      "Cost: \t\t\t$0.0001\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# 1. Single Query Monitoring\n",
    "print(\"1. Single Query Monitoring\")\n",
    "print(\"-\" * 40)\n",
    "\n",
    "with get_openai_callback() as cb:\n",
    "    response = llm.invoke(\"What is the capital of France?\")\n",
    "    print(f\"Response: {response.content}\")\n",
    "    print(\"-\" * 40)\n",
    "    print(f\"Token Usage Statistics:\")\n",
    "    print(f\"Prompt Tokens: \\t\\t{cb.prompt_tokens}\")\n",
    "    print(f\"Completion Tokens: \\t{cb.completion_tokens}\")\n",
    "    print(f\"Total Tokens: \\t\\t{cb.total_tokens}\")\n",
    "    print(f\"Cost: \\t\\t\\t${cb.total_cost:.4f}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2. Multiple Queries Monitoring\n",
      "----------------------------------------\n",
      "Response 1: Python is a high-level, interpreted programming language known for its readability, simplicity, and ...\n",
      "----------------------------------------\n",
      "Response 2: JavaScript is a high-level, dynamic, untyped, and interpreted programming language that is widely us...\n",
      "----------------------------------------\n",
      "Cumulative Statistics:\n",
      "Total Prompt Tokens: \t\t23\n",
      "Total Completion Tokens: \t596\n",
      "Total Tokens: \t\t\t619\n",
      "Total Cost: \t\t\t$0.0060\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# 2. Multiple Queries Monitoring\n",
    "print(\"2. Multiple Queries Monitoring\")\n",
    "print(\"-\" * 40)\n",
    "\n",
    "with get_openai_callback() as cb:\n",
    "    # First query\n",
    "    response1 = llm.invoke(\"What is Python?\")\n",
    "    # Second query\n",
    "    response2 = llm.invoke(\"What is JavaScript?\")\n",
    "\n",
    "    print(f\"Response 1: {response1.content[:100]}...\")\n",
    "    print(\"-\" * 40)\n",
    "    print(f\"Response 2: {response2.content[:100]}...\")\n",
    "    print(\"-\" * 40)\n",
    "    print(\"Cumulative Statistics:\")\n",
    "    print(f\"Total Prompt Tokens: \\t\\t{cb.prompt_tokens}\")\n",
    "    print(f\"Total Completion Tokens: \\t{cb.completion_tokens}\")\n",
    "    print(f\"Total Tokens: \\t\\t\\t{cb.total_tokens}\")\n",
    "    print(f\"Total Cost: \\t\\t\\t${cb.total_cost:.4f}\\n\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain-opentutorial-BXw0bE1H-py3.11",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
